'''
 # @ Author: Jeff Chen
 # @ Create Time: 2025-10-13 15:43:06
 # @ Description: 可视化工具
 '''
import cv2
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, Any

class ThermalVisualizer:
    """热成像可视化工具类"""
    
    def __init__(self, config: Dict[str, Any] = None):
        self.config = config or {}
        
    def create_visualization(self, results: Dict[str, Any], 
                           save_path: str = None) -> np.ndarray:
        """
        创建可视化结果
        
        Args:
            results: 处理结果
            save_path: 保存路径
            
        Returns:
            可视化图像
        """
        original = results['original_image']
        mask = results['low_temp_mask']
        
        # 创建标注图像
        annotated_image = original.copy()
        
        # 在低温区域绘制轮廓
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        vis_config = self.config.get('visualization', {})
        color = vis_config.get('low_temp_color', [0, 0, 255])
        thickness = vis_config.get('contour_thickness', 2)
        
        cv2.drawContours(annotated_image, contours, -1, color, thickness)
        
        # 添加文本信息
        cluster_info = results.get('cluster_info', {})
        if cluster_info:
            low_temp_value = cluster_info.get('low_temperature_value', 0)
            text = f"Low Temp: {low_temp_value:.2f}"
            cv2.putText(annotated_image, text, (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        
        # 保存结果
        if save_path and vis_config.get('save_results', True):
            cv2.imwrite(save_path, annotated_image)
            print(f"结果已保存至: {save_path}")
        
        return annotated_image
    
    def create_comparison_figure(self, results: Dict[str, Any], 
                               save_path: str = None) -> None:
        """
        创建对比图
        
        Args:
            results: 处理结果
            save_path: 保存路径
        """
        fig, axes = plt.subplots(2, 2, figsize=(12, 10))
        
        # 原始图像
        axes[0, 0].imshow(cv2.cvtColor(results['original_image'], cv2.COLOR_BGR2RGB))
        axes[0, 0].set_title('Original Image')
        axes[0, 0].axis('off')
        
        # 灰度图像
        axes[0, 1].imshow(results['gray_image'], cmap='hot')
        axes[0, 1].set_title('Temperature Distribution')
        axes[0, 1].axis('off')
        
        # 低温区域掩码
        axes[1, 0].imshow(results['low_temp_mask'], cmap='gray')
        axes[1, 0].set_title('Low Temperature Mask')
        axes[1, 0].axis('off')
        
        # 标注结果
        annotated = self.create_visualization(results)
        axes[1, 1].imshow(cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB))
        axes[1, 1].set_title('Annotated Result')
        axes[1, 1].axis('off')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"对比图已保存至: {save_path}")
        
        plt.show()
